#!/bin/bash

# 设置变量
AWS_REGION="ap-northeast-1"
AWS_ACCOUNT_ID=$(aws sts get-caller-identity --query "Account" --output text)
POLICY_NAME="AmazonEKSEFSCSIPolicy"

# 获取 EKS 集群列表
CLUSTERS=($(aws eks list-clusters --region ${AWS_REGION} --query 'clusters[]' --output text))
CLUSTER_COUNT=${#CLUSTERS[@]}

if [ $CLUSTER_COUNT -eq 1 ]; then
    # 只有一个集群，自动使用
    CLUSTER_NAME=${CLUSTERS[0]}
    echo "检测到单个 EKS 集群，自动使用: ${CLUSTER_NAME}"
else
    # 多个集群或没有集群，显示列表并要求选择
    echo "检测到 ${CLUSTER_COUNT} 个 EKS 集群:"
    for i in "${!CLUSTERS[@]}"; do
        echo "[$i] ${CLUSTERS[$i]}"
    done
    read -p "请选择集群序号（或输入完整的集群名称）: " CHOICE
    
    # 检查输入是否为数字且在有效范围内
    if [[ $CHOICE =~ ^[0-9]+$ ]] && [ $CHOICE -lt $CLUSTER_COUNT ]; then
        CLUSTER_NAME=${CLUSTERS[$CHOICE]}
    else
        CLUSTER_NAME=$CHOICE
    fi
fi

echo "使用集群: ${CLUSTER_NAME}"
echo "开始配置 EFS 相关 IAM 角色..."

# 创建 EFS 访问策略
echo "创建 IAM 策略..."
cat <<EOF > efs-policy.json
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "elasticfilesystem:DescribeAccessPoints",
        "elasticfilesystem:DescribeFileSystems",
        "elasticfilesystem:DescribeMountTargets",
        "elasticfilesystem:DescribeMountTargetSecurityGroups",
        "ec2:DescribeAvailabilityZones"
      ],
      "Resource": "*"
    }
  ]
}
EOF

# 创建或更新策略
POLICY_ARN="arn:aws:iam::${AWS_ACCOUNT_ID}:policy/${POLICY_NAME}"
if aws iam get-policy --policy-arn ${POLICY_ARN} 2>/dev/null; then
    echo "策略已存在，创建新版本..."
    # 直接创建新版本作为默认版本
    aws iam create-policy-version \
        --policy-arn ${POLICY_ARN} \
        --policy-document file://efs-policy.json \
        --set-as-default

    # 清理旧版本（保留最新的5个版本中的前4个）
    VERSIONS=$(aws iam list-policy-versions --policy-arn ${POLICY_ARN} --query 'Versions[?IsDefaultVersion==`false`].[VersionId]' --output text)
    COUNT=0
    for VERSION in $VERSIONS; do
        COUNT=$((COUNT + 1))
        if [ $COUNT -gt 4 ]; then
            echo "删除旧版本策略: $VERSION"
            aws iam delete-policy-version --policy-arn ${POLICY_ARN} --version-id $VERSION
        fi
    done
else
    aws iam create-policy \
        --policy-name ${POLICY_NAME} \
        --policy-document file://efs-policy.json \
        --region ${AWS_REGION}
    echo "策略创建成功"
fi

# 获取 OIDC provider URL
echo "获取 OIDC provider..."
OIDC_PROVIDER=$(aws eks describe-cluster --name ${CLUSTER_NAME} --region ${AWS_REGION} --query "cluster.identity.oidc.issuer" --output text | sed -e "s/^https:\/\///")

# 创建 CSI Driver 信任关系
echo "创建 CSI Driver 信任关系策略..."
cat <<EOF > csi-trust-policy.json
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "Federated": "arn:aws:iam::${AWS_ACCOUNT_ID}:oidc-provider/${OIDC_PROVIDER}"
      },
      "Action": "sts:AssumeRoleWithWebIdentity",
      "Condition": {
        "StringEquals": {
          "${OIDC_PROVIDER}:aud": "sts.amazonaws.com",
          "${OIDC_PROVIDER}:sub": "system:serviceaccount:kube-system:efs-csi-controller-sa"
        }
      }
    }
  ]
}
EOF

# 创建或更新 CSI Driver 角色
ROLE_NAME="AmazonEKS_EFS_CSI_DriverRole"
if aws iam get-role --role-name ${ROLE_NAME} 2>/dev/null; then
    echo "CSI Driver 角色已存在，更新信任关系..."
    aws iam update-assume-role-policy \
        --role-name ${ROLE_NAME} \
        --policy-document file://csi-trust-policy.json \
        --region ${AWS_REGION}
else
    echo "创建 CSI Driver 角色..."
    aws iam create-role \
        --role-name ${ROLE_NAME} \
        --assume-role-policy-document file://csi-trust-policy.json \
        --region ${AWS_REGION}
fi

# 附加策略到 CSI Driver 角色
echo "附加策略到 CSI Driver 角色..."
aws iam attach-role-policy \
    --policy-arn ${POLICY_ARN} \
    --role-name ${ROLE_NAME} \
    --region ${AWS_REGION}

# 获取角色 ARN
ROLE_ARN=$(aws iam get-role --role-name ${ROLE_NAME} --region ${AWS_REGION} --query 'Role.Arn' --output text)

# 创建或更新 ServiceAccount
echo "配置 ServiceAccount..."
cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: ServiceAccount
metadata:
  name: efs-csi-controller-sa
  namespace: kube-system
  annotations:
    eks.amazonaws.com/role-arn: ${ROLE_ARN}
EOF

# 重启 EFS CSI Driver pods
echo "重启 EFS CSI Driver pods..."
kubectl delete pods -n kube-system -l app.kubernetes.io/name=aws-efs-csi-driver

echo "清理临时文件..."
rm -f efs-policy.json csi-trust-policy.json

echo "配置完成！"
echo "CSI Driver 角色 ARN: ${ROLE_ARN}"
echo "OIDC Provider: ${OIDC_PROVIDER}"

# 验证配置
echo -e "\n=== 验证配置 ==="
echo "1. 检查 CSI Driver 角色权限:"
aws iam list-attached-role-policies --role-name ${ROLE_NAME} --region ${AWS_REGION}

echo -e "\n2. 检查 ServiceAccount:"
kubectl get serviceaccount efs-csi-controller-sa -n kube-system -o yaml

echo -e "\n3. 检查 EFS CSI Driver pods:"
kubectl get pods -n kube-system -l app.kubernetes.io/name=aws-efs-csi-driver

# 检查 pod 日志
echo -e "\n4. 检查 CSI Driver pod 日志:"
POD_NAME=$(kubectl get pods -n kube-system -l app.kubernetes.io/name=aws-efs-csi-driver -o jsonpath='{.items[0].metadata.name}')
kubectl logs -n kube-system ${POD_NAME} -c efs-plugin
